#include <iostream>
#include <unordered_map>
#include <Utils.h>
#include <pcl/point_cloud.h>
#include <pcl/point_types.h>
#include <pcl/io/pcd_io.h>
#include <pcl/segmentation/sac_segmentation.h>
#include <pcl/ModelCoefficients.h>
#include <pcl/segmentation/extract_clusters.h>
#include <pcl/kdtree/kdtree.h>
#include <pcl/filters/extract_indices.h>
#include "plane_info.h"
using namespace std;
using namespace lyxutils::abbreviations;

typedef pcl::PointXYZ PT;
typedef pcl::PointCloud<PT> CloudT;

Vec3f farthestPt(const CloudT &cloud,const Vec3f &direct){
    float max_len=0;
    const PT &pRef=cloud.points[0];
    Vec3f ref={pRef.x,pRef.y,pRef.z};
    Vec3f farPt=ref;
    for(int i=1;i<cloud.points.size();++i){
        const PT &p=cloud.points[i];
        Vec3f v={p.x,p.y,p.z};
        float len=(v-ref).dot(direct);
        if(len>max_len){
            max_len=len;
            farPt=v;
        }
    }
    return farPt;
}

void printUsage(){
    cout<<"usage:ransac_planes pc.pcd|folder"<<endl;
    cout<<"      pc.pcd  =>input point cloud file,contains at least x,y,z"<<endl;
    cout<<"      folder  =>a folder that contains point cloud files"<<endl;
    cout<<"      this command can process one file or a folder contains multiple files"<<endl;
    cout<<"      if process one file,the output file name will be result_<index>.pcd under target folder;"<<endl;
    cout<<"      if process multiple files,the output file names will be <input_file_name>_<index>.pcd"<<endl;
    cout<<"arguments:"<<endl;
    cout<<"  -h        print this help message"<<endl;
    cout<<"  -d=dis    distance threshold of ransac plane fitting(default is 0.2)"<<endl;
    cout<<"  -m=minPts set minimum points in a ransac plane(default is 200)"<<endl;
    cout<<"  -o=folder set output folder(default is `output` folder under current directory)"<<endl;
    cout<<"  -a=angle  set angle limit(with regard to vertical) to filter out horizontal planes"<<endl;
    cout<<"            range (-oo,90) if angle is negative(default),the feature is disabled"<<endl;
    cout<<"  -c        enable connectivity filter to remove points on the extended(not connected) plane"<<endl;
    cout<<"            the distance tolerance is set as maximum side distance(as set by -g option)"<<endl;
    cout<<"  -r=ratio  set cluster size ratio(default is 0.1,this option is only effective when -c is set)"<<endl;
    cout<<"            when connectivity filter is enabled,the points inside plane will be split into clusters"<<endl;
    cout<<"            using euclidean cluster technique.Clusters will be sorted by point size(in decreasing order)"<<endl;
    cout<<"            and clusters after the one with point size less than ratio*last-cluster-point-size will be removed"<<endl;
    cout<<"  -p        enable merge of parallel and close enough planes"<<endl;
    cout<<"            consider walls with normal vector angle less than 5 degree as parallel"<<endl;
    cout<<"            and vertical distance less than 1m as 'close enough',maximum side distance<gap"<<endl;
    cout<<"  -g=gap    set maximum side gap allowed between two planes to be merged(default is 2)"<<endl;
    cout<<"            note: max side gap is only allowed when vertical distance is zero, and decreases"<<endl;
    cout<<"            as vertical distance increases to max(this option is effective when -p is enabled)"<<endl;
}

class RansacSegmentor{
private:
    string filename;
    float distance;
    int minPts;
    float impossibleX;
    float horizontal_angle;
    bool plane_merge;
    bool connectivityFilter;
    float clusterRatio;
    float maxGap;
    CloudT::Ptr cloud;
    Vec3f maxXYZ;
    bool same_plane(const PlaneInfo &plane1,const CloudT &cloud1,
            const PlaneInfo &plane2,const CloudT &cloud2)const{
        //return if two planes are the same plane
        // normal vector angle limit:5 degree;distance limit:1m
        float angle=5.0;
        float dis_limit=1;
        if(acos(abs(plane1.zAxis.dot(plane2.zAxis))/(plane1.zAxis.mold()*plane2.zAxis.mold()))>angle*M_PI/180)return false;
        float dis=abs((plane1.center-plane2.center).dot(plane2.zAxis))/plane2.zAxis.mold();
        if(dis>dis_limit)return false;
        float center_dis=(plane1.center-plane2.center).mold();
        if(center_dis<=dis_limit)return true;//center_dis==0 is included in this situation
        //get farthest point along given direction
        Vec3f direct12=plane2.center-plane1.center;
        direct12=direct12/direct12.mold();//normalize
        Vec3f direct21=-direct12;
        Vec3f farPt1=farthestPt(cloud1,direct12);
        Vec3f farPt2=farthestPt(cloud2,direct21);
        float sum_dis=(farPt1-plane1.center).dot(direct12)+(farPt2-plane2.center).dot(direct21);
        //scale proximity to [-sum_dis,maxGap],sum_dis+proximity is maximum allowed side distance
        float proximity=maxGap*(dis_limit-dis)/dis_limit;
        if(sum_dis+proximity<center_dis)return false;
        return true;
    }
    //计算点云在plane的局部坐标系下的最大最小坐标点,plane的平面参数xAxis,yAxis,zAxis
    void getLocalMinMax(const CloudT &cloud, PlaneInfo &plane)const{
        PT pTemp=cloud.points[0];
        Vec3f pVec={pTemp.x,pTemp.y,pTemp.z};
        Vec3f vecTemp=pVec-plane.center;
        plane.minXYZ={plane.xAxis.dot(vecTemp),plane.yAxis.dot(vecTemp),plane.zAxis.dot(vecTemp)};
        plane.maxXYZ=plane.minXYZ;
        for(int j=0;j<cloud.size();++j){
            const PT &p=cloud.points[j];
            Vec3f pVec=Vec3f({p.x,p.y,p.z})-plane.center;
            pVec={plane.xAxis.dot(pVec),plane.yAxis.dot(pVec),plane.zAxis.dot(pVec)};
            if(pVec.x<plane.minXYZ.x){
                plane.minXYZ.x=pVec.x;
            }
            else if(pVec.x>plane.maxXYZ.x){
                plane.maxXYZ.x=pVec.x;
            }
            if(pVec.y<plane.minXYZ.y){
                plane.minXYZ.y=pVec.y;
            }
            else if(pVec.y>plane.maxXYZ.y){
                plane.maxXYZ.y=pVec.y;
            }
            if(pVec.z<plane.minXYZ.z){
                plane.minXYZ.z=pVec.z;
            }
            else if(pVec.z>plane.maxXYZ.z){
                plane.maxXYZ.z=pVec.z;
            }
        }
    }
    //confirm cloud.width*cloud.height==cloud.points.size()==plane.ptCount
    void confirmPointCount(CloudT &cloud,PlaneInfo &plane)const{
        cloud.height=1;
        cloud.width=cloud.points.size();
        plane.ptCount=cloud.width;
    }
public:
    RansacSegmentor(){
        maxXYZ={FLT_MIN,FLT_MIN,FLT_MIN};
        connectivityFilter=false;
    }
    void setInputCloud(string filename){
        this->filename=xos::abspath(filename);
        cloud=CloudT::Ptr(new CloudT);
        if(xos::suffix(filename)=="pcd"){
            pcl::io::loadPCDFile(filename,*cloud);
        }
        else if(xos::suffix(filename)=="txt"){
            xio::read_pts(filename,cloud->points);
        }
        else{
            throw invalid_argument("invalid argument:input point cloud file type unrecognized,file name:"+filename);
        }
//        pcl::io::loadPCDFile(filename,*cloud);
        for(int i=1;i<cloud->size();++i){
            const PT &p=cloud->points[i];
            if(maxXYZ.x<p.x)maxXYZ.x=p.x;
            if(maxXYZ.y<p.y)maxXYZ.y=p.y;
            if(maxXYZ.z<p.z)maxXYZ.z=p.z;
        }
        impossibleX=maxXYZ.x+1;//set impossibleX as an impossible x value
    }
    void setDistanceThresh(float dis){
        distance=dis;
    }
    void setMinPts(int pts){
        minPts=pts;
    }
    void setHorizontalDamper(float angle){
        horizontal_angle=angle;
    }
    void enablePlaneMerge(bool enabled){
        plane_merge=enabled;
    }
    void setMaxGap(float gap){
        maxGap=gap;
    }
    void setClusterRatio(float ratio){
        clusterRatio=ratio;
    }
    void segment(vector<PlaneInfo> &planes,CloudT::CloudVectorType &plane_clouds){
        int startIdx=planes.size();//planes is accumulative start from specific index
        while(cloud->size()>0){
            pcl::ModelCoefficients::Ptr coefficients(new pcl::ModelCoefficients);
            pcl::PointIndices::Ptr inliers(new pcl::PointIndices);
            pcl::SACSegmentation<pcl::PointXYZ> seg;
            seg.setInputCloud(cloud);
            seg.setOptimizeCoefficients(true);
            seg.setModelType(pcl::SACMODEL_PLANE);
            seg.setMethodType(pcl::SAC_RANSAC);
            seg.setMaxIterations(1000);
            seg.setDistanceThreshold(distance);
            seg.segment(*inliers,*coefficients);
            //如果拟合的平面点数小于最小点数限制,终止循环,因为之后ransac提取到的平面都不会比这个大
            if(inliers->indices.size()<minPts){
                break;
            }
            PlaneInfo plane_tmp;
            planeFit(cloud->points, inliers->indices, plane_tmp.center, plane_tmp.sigma, plane_tmp.xAxis, plane_tmp.yAxis, plane_tmp.zAxis);
            bool noHorizontal=horizontal_angle<0 || horizontal_angle >90 ||
                    acos(abs(plane_tmp.zAxis.dot(Vec3f(0,0,1)))/plane_tmp.zAxis.mold())>horizontal_angle*M_PI/180;
            if(noHorizontal){
                //extract points in a single plane
                plane_clouds.push_back(CloudT());
                pcl::ExtractIndices<PT> extract;
                extract.setInputCloud(cloud);
                extract.setIndices(inliers);
                extract.setNegative(false);
                CloudT &outCloud=plane_clouds[plane_clouds.size()-1];
                extract.filter(outCloud);
                confirmPointCount(outCloud,plane_tmp);
                getLocalMinMax(outCloud,plane_tmp);
                planes.push_back(plane_tmp);
            }
            for(int i=0;i<inliers->indices.size();++i){
//                cloud->points[i].x=impossibleX;//set point x as impossible x for later deletion
                //fix bug:delete points on detected plane,deleted the wrong points in above line
                //this should correct the ransac plane extract procedure to extract all possible planes
                cloud->points[inliers->indices[i]].x=impossibleX;
            }
            for(auto it=cloud->points.begin();it!=cloud->points.end();){//delete plane_tmp points
                if(it->x==impossibleX){
                    it=cloud->points.erase(it);
                }
                else{
                    ++it;
                }
            }
            cloud->height=1;
            cloud->width=cloud->points.size();
        }//end while loop:continue detecting ransac plane when cloud has point inside
        //merge planes and corresponding clouds
        vector<int> need_recalculate(plane_clouds.size(),0);//whether the plane need re-calculating
        if(plane_merge){
            //if no new planes detected,there is no need for merge!
            if(plane_clouds.size()>startIdx){
                need_recalculate.clear();
                for(int i=0;i<plane_clouds.size()-1;++i){
                    //if plane_clouds.size()==0,this will be a dead loop
                    // (because (uint)0-1=MAX_UINT, of course i<MAX_UINT)
                    if(planes[i].ptCount==0)continue;
                    need_recalculate.push_back(0);
                    int start=i+1>startIdx?i+1:startIdx;
                    for(int j=start;j<plane_clouds.size();++j){
                        if(planes[j].ptCount==0)continue;
                        if(same_plane(planes[i],plane_clouds[i],planes[j],plane_clouds[j])){
                            plane_clouds[i]+=plane_clouds[j];
                            confirmPointCount(plane_clouds[i],planes[i]);
                            *(need_recalculate.end()-1)=1;//the i-th plane need recalculation
                            planes[j].ptCount=0;    //set to zero to delete the plane & cloud
                            plane_clouds[j].width=0;//width is uint and thus cannot be negative
                        }
                    }
                }
            }
            //delete the redundant
            for(auto it=planes.begin()+startIdx;it!=planes.end();){
                if(it->ptCount==0){
                    it=planes.erase(it);
                }
                else{
                    ++it;
                }
            }
            for(auto it=plane_clouds.begin()+startIdx;it!=plane_clouds.end();){
                if(it->width==0){
                    it=plane_clouds.erase(it);
                }
                else{
                    ++it;
                }
            }
        }
        //recalculate plane infomation
        int count=0;
        for(int i=0;i<plane_clouds.size();++i){
            PlaneInfo &plane_temp=planes[i];
            if(need_recalculate[i]){
                vector<int> indices;
                np::range<int>(indices,0,plane_clouds[i].size());
                //least square fit plane
                planeFit(plane_clouds[i].points,indices,plane_temp.center,plane_temp.sigma,
                        plane_temp.xAxis,plane_temp.yAxis,plane_temp.zAxis);
                //get min max coordinates in each plane
                getLocalMinMax(plane_clouds[i],plane_temp);
                confirmPointCount(plane_clouds[i],plane_temp);
            }
        }
    }//end of function segment
    void optimisePlane(PlaneInfo &plane,CloudT &cloud){
        vector<pcl::PointIndices> cluster_indices;
        pcl::EuclideanClusterExtraction<PT> ec;
        pcl::search::KdTree<PT>::Ptr tree(new pcl::search::KdTree<PT>());
        CloudT::Ptr cloudTemp(new CloudT());
        *cloudTemp+=cloud;
        tree->setInputCloud(cloudTemp);
        ec.setInputCloud(cloudTemp);
        ec.setMinClusterSize(1);
        ec.setMaxClusterSize(cloudTemp->points.size());
        ec.setSearchMethod(tree);
        ec.setClusterTolerance(maxGap);
        ec.extract(cluster_indices);
        pcl::ExtractIndices<PT> extractor;
        extractor.setInputCloud(cloudTemp);
        extractor.setNegative(false);
        if(cluster_indices.size()==1){//only one cluster
            return;
        }
        else{
            cloud.points.clear();//clear point cloud
            int minSz=max(minPts,(int)(cluster_indices[0].indices.size()*clusterRatio));
            //select cluster with point size no less than minPts and ratio*largest_cluster_size
            for(int j=0;j<cluster_indices.size()&&cluster_indices[j].indices.size()>=minSz;++j){
                CloudT cloudSingle;
                extractor.setIndices(boost::make_shared<vector<int> >(cluster_indices[j].indices));
                extractor.filter(cloudSingle);
                cloud+=cloudSingle;
//                if(cluster_indices[j+1].indices.size()<clusterRatio*cluster_indices[j].indices.size()){
//                    break;//break loop when next cluster size is way too small compared to the current one
//                }
            }
        }
        //if number of points on plane is changed=>need to recalculate the plane parameters
        if(cloud.points.size()!=cloudTemp->points.size()){
            vector<int> indices;
            np::range<int>(indices,0,cloud.points.size());
            //least square fit plane
            planeFit(cloud.points,indices,plane.center,plane.sigma,
                     plane.xAxis,plane.yAxis,plane.zAxis);
            //get min max coordinates in each plane
            getLocalMinMax(cloud,plane);
            confirmPointCount(cloud,plane);
        }
    }
    Vec3f getMaxPos(){
        return maxXYZ;
    }
};

int main(int argc,char **argv) {
    xio::CLParser parser;
    parser.parse(argc,argv);
    vector<string> paras=parser.getParameters();
    if(paras.size()!=1 || parser.hasOption("h")){
        printUsage();
        return 1;
    }
    //deal with options
    float distance=parser.hasOption("d")?stof(parser.getOptionArg("d")):0.2;
    int minPts=parser.hasOption("m")?stoi(parser.getOptionArg("m")):200;
    float angle=parser.hasOption("a")?stof(parser.getOptionArg("a")):-1;
    float gap=parser.hasOption("g")?stof(parser.getOptionArg("g")):2;
    string outFolder=xos::abspath(parser.hasOption("o")?parser.getOptionArg("o"):"output");
    xos::mkdir(outFolder);
    //segment file(s) and get planes
    RansacSegmentor segmentor;
    segmentor.setDistanceThresh(distance);
    segmentor.setMinPts(minPts);
    segmentor.setHorizontalDamper(angle);
    segmentor.enablePlaneMerge(parser.hasOption("p"));
    segmentor.setMaxGap(gap);
    cout<<"processing..."<<endl;
    vector<PlaneInfo> planes;//plane infomation
    CloudT::CloudVectorType plane_clouds;
    if(xos::directory(paras[0])){//process a folder
        vector<string> files=xos::lsfile(paras[0]);
        std::sort(files.begin(),files.end(),[](string a,string b)->bool{
            return a.size()<b.size() || (a.size()==b.size() && a<b);
        });
        int count=1;
        for(string f:files){
//            if(xos::suffix(f)!="pcd")continue;
            segmentor.setInputCloud(f);
            segmentor.segment(planes,plane_clouds);
            xio::consoleProgressBar(count++,files.size());
        }
        if(count<=files.size()){
            xio::consoleProgressBar(100,100);
        }
    }
    else if(xos::file(paras[0])){//process single file
//        if(xos::suffix(paras[0])=="pcd"){
//            segmentor.setInputCloud(paras[0]);
//            segmentor.segment(planes,plane_clouds);
//        }
        segmentor.setInputCloud(paras[0]);
        segmentor.segment(planes,plane_clouds);
    }
    else{
        cout<<"file:"<<paras[0]<<" not exist!"<<endl;
        return 1;
    }
    //optimise plane to remove non-connected parts of plane
    if(parser.hasOption("c")){
        float ratio=parser.hasOption("r")?stof(parser.getOptionArg("r")):0.1;
        segmentor.setClusterRatio(ratio);
        for(int i=0;i<planes.size();++i){
            segmentor.optimisePlane(planes[i],plane_clouds[i]);
        }
        //remove empty plane
        auto it=planes.begin();
        auto c_it=plane_clouds.begin();
        for(;it!=planes.end()&&c_it!=plane_clouds.end();){
            if(c_it->size()==0){
                it=planes.erase(it);
                c_it=plane_clouds.erase(c_it);
            }
            else{
                it++;
                c_it++;
            }
        }
    }
    Vec3f viewPos=segmentor.getMaxPos()+Vec3f({1,1,1});
    for(int i=0;i<planes.size();++i){
        planes[i].setViewPoint(viewPos);
    }
    //save plane info to file
    ofstream ofs(outFolder+"/planes.txt");
    for(int i=0;i<planes.size();++i){
        ofs<<planes[i];
    }
    ofs.flush();
    ofs.close();
    //save point cloud of plane to file
    int idx=0;
    for(const auto &cloud:plane_clouds){
        pcl::io::savePCDFile(outFolder+"/result_plane_"+str::type2str(idx++)+".pcd",cloud);
    }
    return 0;
}